import logging
from elasticsearch import Elasticsearch
from typing import List


class ESSearch:
    def __init__(
        self,
        host: str,
        user: str,
        password: str,
        verify_certs=False,
    ):
        """
        创建 Elasticsearch 连接
        :param es_host: Elasticsearch 主机地址
        :param es_auth: 认证信息
        :param verify_certs: 是否验证服务器的 TLS 证书,默认为 False
        :return: Elasticsearch 连接对象
        """
        logging.info("正在创建 Elasticsearch 连接...")
        if host[:6] == "https":
            import ssl

            self.es = Elasticsearch(
                [host],
                basic_auth=(user, password),
                ssl_context=(
                    None if verify_certs else ssl._create_unverified_context()
                ),
                verify_certs=verify_certs,
            )
        else:
            self.es = Elasticsearch(
                [host],
                basic_auth=(user, password),
            )

        logging.info("Elasticsearch 连接创建成功")

    def query(self, query: str, index_name="navy_poc_20250306", min_length=1):
        es_result = self.es_ikmax_slangs_function_score(
            index_name, query, min_length=min_length
        )
        return es_result

    def es_ikmax_slangs_function_score(
        self, index_name: str, query: str, min_length: int
    ) -> List[str]:

        # 如果没有提取到 slang，只使用 content.chinese 匹配
        chunk_body = {
            "query": {
                "match": {
                    "content.chinese": {
                        "query": query,
                        "analyzer": "ik_max_word",
                    }
                }
            },
            "_source": ["content", "file_name"],
            "size": 5,
        }
        try:
            chunk_result: dict = self.es.search(index=index_name, body=chunk_body)
            if chunk_result["hits"]["hits"]:
                chunks_content: list[str] = [
                    hit["_source"]["content"] for hit in chunk_result["hits"]["hits"]
                ]
                chunks_source: list[str] = [
                    hit["_source"]["file_name"] for hit in chunk_result["hits"]["hits"]
                ]
                return chunks_content, chunks_source
            else:
                logging.error("No chunks found matching the query")
                return [], []
        except Exception as e:
            logging.error(f"Error in Elasticsearch search: {e}")
            return [], []


import yaml

if __name__ == "__main__":
    configs = yaml.safe_load(open("/root/data/zhonghe/config.yaml"))
    es = ESSearch(
        host=configs.get("ES_HOST"),
        user=configs.get("ES_USER"),
        password=configs.get("ES_PASSWORD"),
    )
    print(
        es.query(
            "双枝不等重大型树状结构斜向原位提升施工工法",
            index_name="rag_zhonghe_v2_20250102",
            min_length=1,
        )
    )
